#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PaddleOCR Flask Web服务
提供RESTful API接口供Django调用
"""

import os
import time
import json
from pathlib import Path
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
import tempfile

from paddleocr import PaddleOCR

# 全局OCR对象
ocr = None

def init_ocr():
    """初始化OCR模型"""
    global ocr
    if ocr is None:
        print("正在初始化PaddleOCR模型...")
        ocr = PaddleOCR(
            use_doc_orientation_classify=True,
            use_doc_unwarping=True
        )
        print("PaddleOCR模型初始化完成！")
    return ocr

# 创建Flask应用
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传文件大小为16MB

# 支持的图片格式
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'tif', 'webp'}

def allowed_file(filename):
    """检查文件格式是否支持"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/api/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        'status': 'healthy',
        'message': 'OCR服务运行正常',
        'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
    })

@app.route('/api/ocr/recognize', methods=['POST'])
def recognize_text():
    """OCR文字识别接口"""
    try:
        # 检查是否有文件上传
        if 'image' not in request.files:
            return jsonify({
                'success': False,
                'error': '没有上传图片文件'
            }), 400
        
        file = request.files['image']
        
        # 检查文件名
        if file.filename == '':
            return jsonify({
                'success': False,
                'error': '没有选择文件'
            }), 400
        
        # 检查文件格式
        if not allowed_file(file.filename):
            return jsonify({
                'success': False,
                'error': f'不支持的文件格式，支持的格式: {", ".join(ALLOWED_EXTENSIONS)}'
            }), 400
        
        # 获取OCR实例
        ocr_instance = init_ocr()
        
        # 记录开始时间
        start_time = time.time()
        
        # 保存临时文件
        filename = secure_filename(file.filename)
        temp_path = os.path.join(tempfile.gettempdir(), filename)
        file.save(temp_path)
        
        try:
            # 执行OCR识别
            result = ocr_instance.predict(temp_path)
            
            # 记录结束时间
            end_time = time.time()
            processing_time = end_time - start_time
            
            # 提取识别结果
            if result and len(result) > 0:
                res = result[0]
                
                # 从OCRResult对象的json属性中提取数据
                json_data = res.json
                
                if 'res' in json_data:
                    res_data = json_data['res']
                    
                    if isinstance(res_data, dict):
                        rec_texts = res_data.get('rec_texts', [])
                        rec_scores = res_data.get('rec_scores', [])
                        
                        if rec_texts:
                            # 计算平均置信度
                            avg_confidence = sum(rec_scores) / len(rec_scores) if rec_scores else 0.0
                            
                            # 合并所有文本
                            combined_text = '\n'.join(rec_texts)
                            
                            # 返回结果
                            return jsonify({
                                'success': True,
                                'data': {
                                    'filename': filename,
                                    'processing_time': round(processing_time, 2),
                                    'avg_confidence': round(avg_confidence, 4),
                                    'text_count': len(rec_texts),
                                    'texts': rec_texts,
                                    'combined_text': combined_text,
                                    'scores': rec_scores,
                                    'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
                                }
                            })
                        else:
                            return jsonify({
                                'success': True,
                                'data': {
                                    'filename': filename,
                                    'processing_time': round(processing_time, 2),
                                    'avg_confidence': 0.0,
                                    'text_count': 0,
                                    'texts': [],
                                    'combined_text': '',
                                    'scores': [],
                                    'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
                                    'message': '未识别到文字'
                                }
                            })
                    else:
                        return jsonify({
                            'success': False,
                            'error': 'OCR结果数据格式错误'
                        }), 500
                else:
                    return jsonify({
                        'success': False,
                        'error': 'OCR结果中未找到数据'
                    }), 500
            else:
                return jsonify({
                    'success': False,
                    'error': 'OCR识别失败，未获得结果'
                }), 500
                
        finally:
            # 清理临时文件
            if os.path.exists(temp_path):
                os.remove(temp_path)
                
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'OCR处理失败: {str(e)}'
        }), 500

@app.route('/api/ocr/batch', methods=['POST'])
def batch_recognize():
    """批量OCR识别接口"""
    try:
        # 检查是否有文件上传
        if 'images' not in request.files:
            return jsonify({
                'success': False,
                'error': '没有上传图片文件'
            }), 400
        
        files = request.files.getlist('images')
        
        if not files:
            return jsonify({
                'success': False,
                'error': '没有选择文件'
            }), 400
        
        # 获取OCR实例
        ocr_instance = init_ocr()
        
        results = []
        total_start_time = time.time()
        
        for file in files:
            if file.filename == '':
                continue
                
            if not allowed_file(file.filename):
                results.append({
                    'filename': file.filename,
                    'success': False,
                    'error': f'不支持的文件格式'
                })
                continue
            
            try:
                # 记录开始时间
                start_time = time.time()
                
                # 保存临时文件
                filename = secure_filename(file.filename)
                temp_path = os.path.join(tempfile.gettempdir(), filename)
                file.save(temp_path)
                
                try:
                    # 执行OCR识别
                    result = ocr_instance.predict(temp_path)
                    
                    # 记录结束时间
                    end_time = time.time()
                    processing_time = end_time - start_time
                    
                    # 提取识别结果
                    if result and len(result) > 0:
                        res = result[0]
                        json_data = res.json
                        
                        if 'res' in json_data:
                            res_data = json_data['res']
                            
                            if isinstance(res_data, dict):
                                rec_texts = res_data.get('rec_texts', [])
                                rec_scores = res_data.get('rec_scores', [])
                                
                                avg_confidence = sum(rec_scores) / len(rec_scores) if rec_scores else 0.0
                                combined_text = '\n'.join(rec_texts)
                                
                                results.append({
                                    'filename': filename,
                                    'success': True,
                                    'processing_time': round(processing_time, 2),
                                    'avg_confidence': round(avg_confidence, 4),
                                    'text_count': len(rec_texts),
                                    'texts': rec_texts,
                                    'combined_text': combined_text,
                                    'scores': rec_scores
                                })
                            else:
                                results.append({
                                    'filename': filename,
                                    'success': False,
                                    'error': 'OCR结果数据格式错误'
                                })
                        else:
                            results.append({
                                'filename': filename,
                                'success': False,
                                'error': 'OCR结果中未找到数据'
                            })
                    else:
                        results.append({
                            'filename': filename,
                            'success': False,
                            'error': 'OCR识别失败'
                        })
                        
                finally:
                    # 清理临时文件
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
                        
            except Exception as e:
                results.append({
                    'filename': file.filename,
                    'success': False,
                    'error': f'处理失败: {str(e)}'
                })
        
        total_time = time.time() - total_start_time
        
        return jsonify({
            'success': True,
            'data': {
                'total_files': len(files),
                'total_time': round(total_time, 2),
                'results': results,
                'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
            }
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'批量处理失败: {str(e)}'
        }), 500

if __name__ == '__main__':
    print("正在启动PaddleOCR Flask服务...")
    print("服务地址: http://localhost:5000")
    print("API文档:")
    print("  GET  /api/health - 健康检查")
    print("  POST /api/ocr/recognize - 单张图片识别")
    print("  POST /api/ocr/batch - 批量图片识别")
    print("="*50)
    
    # 首次初始化OCR模型
    init_ocr()
    
    # 启动Flask服务 - 监听所有网络接口
    print("OCR服务配置:")
    print(f"  服务器IP: 192.168.0.169")
    print(f"  服务端口: 5000")
    print(f"  服务地址: http://192.168.0.169:5000")
    print("="*50)
    
    app.run(host='0.0.0.0', port=5000, debug=False)
